#!/bin/bash
# EbbFlow 本地测试脚本
# 参数: N=4, batchsize=1000, epoch=10

set -e

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

# 配置参数
N_NODES=4
EPOCHS=10
BATCH_SIZE=1000
R=5
FAULTY_NODES=1
PROJECT_ROOT="/home/ebbflow"
LOG_DIR="${PROJECT_ROOT}/logs_test"
BINARY_NAME="ebbflow_server"
BINARY="${PROJECT_ROOT}/${BINARY_NAME}"

# 日志函数
log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

log_step() {
    echo -e "${BLUE}==>${NC} ${GREEN}$1${NC}"
}

# 清理函数
cleanup() {
    log_step "清理进程和日志..."
    pkill -f "$BINARY_NAME" 2>/dev/null || true
    sleep 1
    rm -rf "$LOG_DIR"
    mkdir -p "$LOG_DIR"
    log_info "清理完成"
}

# 编译
build() {
    log_step "编译 EbbFlow Server..."
    cd "$PROJECT_ROOT"
    go build -o "$BINARY_NAME" ./src/main/ebbflow_server.go
    if [ $? -ne 0 ]; then
        log_error "编译失败"
        exit 1
    fi
    log_info "编译成功: $BINARY"
}

# 检查配置
check_config() {
    log_step "检查配置文件..."
    
    # 检查并创建hosts.json
    if [ ! -f "$PROJECT_ROOT/hosts.json" ]; then
        log_info "创建 hosts.json..."
        cat > "$PROJECT_ROOT/hosts.json" << EOF
[
    "localhost:11000",
    "localhost:11001",
    "localhost:11002",
    "localhost:11003"
]
EOF
        log_info "已创建默认hosts.json"
    fi
    
    # 检查配置文件中的batchsize
    if [ -f "$PROJECT_ROOT/etc/conf.json" ]; then
        current_batch=$(grep -o '"maxBatchSize": [0-9]*' "$PROJECT_ROOT/etc/conf.json" | grep -o '[0-9]*')
        if [ "$current_batch" != "$BATCH_SIZE" ]; then
            log_info "更新配置文件中的 maxBatchSize: $current_batch -> $BATCH_SIZE"
            # 使用sed更新batch size
            sed -i "s/\"maxBatchSize\": [0-9]*/\"maxBatchSize\": $BATCH_SIZE/" "$PROJECT_ROOT/etc/conf.json"
        fi
    fi
    
    log_info "配置验证通过: N=$N_NODES, Epochs=$EPOCHS, BatchSize=$BATCH_SIZE, R=$R, F=$FAULTY_NODES"
}

# 启动节点
start_nodes() {
    log_step "启动 $N_NODES 个节点..."
    
    for i in $(seq 0 $((N_NODES - 1))); do
        log_info "启动节点 $i..."
        # 不输出到文件，直接输出到控制台（重定向到日志目录以便查看）
        nohup "$BINARY" -id $i -n $N_NODES -f $FAULTY_NODES -epochs $EPOCHS -r $R > "$LOG_DIR/node_${i}.log" 2>&1 &
        echo $! > "$LOG_DIR/node_${i}.pid"
        sleep 0.5
    done
    
    sleep 2
    log_info "所有节点已启动"
}

# 检查节点状态
check_nodes() {
    log_step "检查节点状态..."
    
    local running=0
    for i in $(seq 0 $((N_NODES - 1))); do
        if pgrep -f "ebbflow_server.*-id $i" > /dev/null; then
            running=$((running + 1))
            log_info "节点 $i: 运行中"
        else
            log_error "节点 $i: 未运行"
        fi
    done
    
    if [ $running -eq $N_NODES ]; then
        log_info "所有 $N_NODES 个节点都在运行"
        return 0
    else
        log_error "只有 $running/$N_NODES 个节点在运行"
        return 1
    fi
}

# 等待测试完成
wait_for_completion() {
    log_step "等待测试完成（epoch=$EPOCHS）..."
    
    # 估算运行时间：每个epoch大约需要几秒，加上一些缓冲
    ESTIMATED_TIME=$((EPOCHS * 3 + 10))
    log_info "预计运行时间: ${ESTIMATED_TIME} 秒"
    
    # 监控日志，查找epoch完成信息
    local completed=0
    local check_count=0
    local max_checks=$((ESTIMATED_TIME * 2))  # 每0.5秒检查一次
    
    while [ $check_count -lt $max_checks ]; do
        # 检查节点0的日志，查找epoch完成信息
        if [ -f "$LOG_DIR/node_0.log" ]; then
            # 查找是否有"epoch"相关的完成信息
            epoch_count=$(grep -i "epoch.*completed\|epoch.*完成\|reached target epoch\|epochs completed" "$LOG_DIR/node_0.log" 2>/dev/null | wc -l || echo "0")
            
            # 查找明确的终止信息
            if grep -q "Reached target epoch\|all.*epochs completed\|epochs completed successfully" "$LOG_DIR/node_0.log" 2>/dev/null; then
                log_info "检测到epoch完成信息"
                sleep 2  # 等待所有节点完成
                break
            fi
        fi
        
        # 检查所有节点是否还在运行
        running_count=0
        for i in $(seq 0 $((N_NODES - 1))); do
            if pgrep -f "ebbflow_server.*-id $i" > /dev/null; then
                running_count=$((running_count + 1))
            fi
        done
        
        # 如果所有节点都停止了，可能是完成了或出错了
        if [ $running_count -eq 0 ]; then
            log_info "所有节点已停止"
            break
        fi
        
        sleep 0.5
        check_count=$((check_count + 1))
        
        # 每10秒显示一次进度
        if [ $((check_count % 20)) -eq 0 ]; then
            log_info "等待中... (已等待 $((check_count / 2)) 秒)"
        fi
    done
    
    if [ $check_count -ge $max_checks ]; then
        log_error "等待超时"
        return 1
    fi
    
    return 0
}

# 显示结果
show_results() {
    log_step "测试结果..."
    
    for i in $(seq 0 $((N_NODES - 1))); do
        if [ -f "$LOG_DIR/node_${i}.log" ]; then
            log_info "=== 节点 $i 日志 (最后30行) ==="
            tail -30 "$LOG_DIR/node_${i}.log" | grep -E "epoch|Epoch|EPOCH|completed|完成|terminated|终止" || tail -30 "$LOG_DIR/node_${i}.log"
            echo ""
        fi
    done
    
    # 检查是否成功达到指定epoch数
    if [ -f "$LOG_DIR/node_0.log" ]; then
        final_epoch=$(grep -i "epoch.*completed\|epoch.*完成" "$LOG_DIR/node_0.log" 2>/dev/null | tail -1)
        if echo "$final_epoch" | grep -q "epoch.*$EPOCHS\|$EPOCHS.*epoch"; then
            log_info "✓ 测试成功：达到了指定的epoch数量 ($EPOCHS)"
        else
            log_info "最终epoch信息: $final_epoch"
        fi
    fi
}

# 主程序
main() {
    log_step "开始本地测试 - N=$N_NODES, BatchSize=$BATCH_SIZE, Epochs=$EPOCHS"
    
    # 清理
    cleanup
    
    # 编译
    build
    
    # 检查配置
    check_config
    
    # 启动节点
    start_nodes
    
    # 检查节点状态
    if ! check_nodes; then
        log_error "节点启动失败"
        exit 1
    fi
    
    # 等待完成
    wait_for_completion
    
    # 显示结果
    show_results
    
    log_step "测试完成"
    log_info "日志文件位于: $LOG_DIR"
    log_info "使用以下命令查看详细日志:"
    echo "  tail -f $LOG_DIR/node_0.log"
}

# 捕获中断信号
trap 'cleanup; exit' INT TERM

# 运行主程序
main


